%%%-------------------------------------------------------------------
%%% @author Viktor Green <viktor@localhost>
%%% @copyright (C) 2011, Viktor Green
%%% @doc
%%%
%%% @end
%%% Created : 14 Dec 2011 by Viktor Green <viktor@localhost>
%%%-------------------------------------------------------------------
-module(torrent_ctrl).

-behaviour(gen_server).

%% API
-export([start_link/3, tracker_callback/2]).

%% gen_server callbacks
-export([init/1, handle_call/3, handle_cast/2, handle_info/2,
	 terminate/2, code_change/3]).

-define(SERVER, ?MODULE).

-include("records.hrl"). 

-record(state, {fileHandler :: pid() | undefined
	       ,torrent :: #torrent{} | undefined
	       ,supervisor :: pid()
	       ,tracker :: pid() | undefined
	       ,peers :: ordset:ordset(tracker:address())
	       ,need_data = false :: boolean()
	       ,pieces = []:: [boolean()]
	       }).

%%%===================================================================
%%% API
%%%===================================================================

%%--------------------------------------------------------------------
%% @doc
%% Starts the server
%%
%% @end
%%--------------------------------------------------------------------
-spec start_link(Torrent :: #torrent{},
		 FileExists :: boolean(), Supervisor :: pid()) ->
			{ok, Pid :: pid()} |
			ignore | 
			{error, Error :: term()}.
start_link(Torrent, FileExists, Supervisor) ->
    gen_server:start_link(?MODULE, [Torrent,
				    FileExists, Supervisor], []).

%%--------------------------------------------------------------------
%% @doc
%% Called from the tracker_srv module, after a connection is done
%% @end
%%--------------------------------------------------------------------
-spec tracker_callback(Pid :: pid(), Peers :: [tracker:address()]) ->
			      ok.
tracker_callback(Pid, Peers) ->
    gen_server:cast(Pid, {tracker_connection, Peers}).

%%%===================================================================
%%% gen_server callbacks
%%%===================================================================

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Initializes the server
%%
%% @spec init(Args) -> {ok, State} |
%%                     {ok, State, Timeout} |
%%                     ignore |
%%                     {stop, Reason}
%% @end
%%--------------------------------------------------------------------
init([#torrent{type = file} = Torrent, FileExists, Supervisor]) ->
    timer:send_after(2000, {do_initial_things, FileExists}),
    {ok, #state{supervisor = Supervisor, peers = ordsets:new(),
	        torrent = Torrent}};
init([#torrent{type = Type}, _, _]) ->
    {stop, {not_implemented, Type}}.


%%--------------------------------------------------------------------
%% @private
%% @doc
%% Handling call messages
%%
%% @spec handle_call(Request, From, State) ->
%%                                   {reply, Reply, State} |
%%                                   {reply, Reply, State, Timeout} |
%%                                   {noreply, State} |
%%                                   {noreply, State, Timeout} |
%%                                   {stop, Reason, Reply, State} |
%%                                   {stop, Reason, State}
%% @end
%%--------------------------------------------------------------------
handle_call(_Request, _From, State) ->
    Reply = ok,
    {reply, Reply, State}.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Handling cast messages
%%
%% @spec handle_cast(Msg, State) -> {noreply, State} |
%%                                  {noreply, State, Timeout} |
%%                                  {stop, Reason, State}
%% @end
%%--------------------------------------------------------------------
handle_cast({tracker_connection, Peers}, State) ->
    NewState = tracker_connection(Peers, State),
    {noreply, NewState};
handle_cast(_Msg, State) ->
    {noreply, State}.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Handling all non call/cast messages
%%
%% @spec handle_info(Info, State) -> {noreply, State} |
%%                                   {noreply, State, Timeout} |
%%                                   {stop, Reason, State}
%% @end
%%--------------------------------------------------------------------
handle_info({do_initial_things, FileExists}, State) ->
    {noreply, initial_things(FileExists, State)};
handle_info(_Info, State) ->
    {noreply, State}.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% This function is called by a gen_server when it is about to
%% terminate. It should be the opposite of Module:init/1 and do any
%% necessary cleaning up. When it returns, the gen_server terminates
%% with Reason. The return value is ignored.
%%
%% @spec terminate(Reason, State) -> void()
%% @end
%%--------------------------------------------------------------------
terminate(_Reason, _State) ->
    ok.

%%--------------------------------------------------------------------
%% @private
%% @doc
%% Convert process state when code is changed
%%
%% @spec code_change(OldVsn, State, Extra) -> {ok, NewState}
%% @end
%%--------------------------------------------------------------------
code_change(_OldVsn, State, _Extra) ->
    {ok, State}.

%%%===================================================================
%%% Internal functions
%%%===================================================================

%%--------------------------------------------------------------------
%% @doc
%% Do some initial things that cannot be done in init (since I'm not
%% sure wheter the FileHandler gen_server has started or not at that
%% point.
%% @end
%%--------------------------------------------------------------------
-spec initial_things(FileExists :: boolean(), State :: #state{}) ->
			    NewState :: #state{}.
initial_things(FileExists, #state{torrent = Torrent,
				  supervisor = Sup} = State) ->
    PL = Torrent#torrent.pieceLength,
    Length = Torrent#torrent.length,
    Pieces = Torrent#torrent.checkSums,
    FH = torrent_sup:get_filehandler(Sup),
    {Left, Haves} = case FileExists of
			 true -> 
			     Have = check_file(FH, PL, 0, Length,
					       Pieces, []),
			     {how_much_left(Have, Length, PL), Have};
			 false ->
			     file_handler:allocate(FH, Length),
			     {0, generate_empty_have(Length, PL)}
		   end,
    NeedData = case Left of
		   0 -> false;
		   _ -> true
	       end,
    {ok, Track} = torrent_sup:start_tracking(Sup, Torrent,
					     Left, self()),
    State#state{fileHandler = FH, tracker = Track,
		need_data = NeedData, pieces = Haves}.

%%--------------------------------------------------------------------
%% @doc
%% Calculate how much data is left to download
%% @end
%%--------------------------------------------------------------------
-spec how_much_left(HavePieces :: [boolean()], Length :: integer(),
		    PieceLength :: integer()) ->
			   integer().
how_much_left([true | Tail], Length, PieceLength) ->
    how_much_left(Tail, Length, PieceLength);
how_much_left([false | []], Length, PieceLength) ->
    Length rem PieceLength;
how_much_left([false | Tail], Length, PieceLength) ->
    PieceLength + how_much_left(Tail, Length, PieceLength).

%%--------------------------------------------------------------------
%% @doc
%% Validate file
%% @end
%%--------------------------------------------------------------------
-spec check_file(FileHandler :: pid(), PieceLength :: integer(),
		 ID :: integer(), Length :: integer(),
		 PieceChecks :: string(), Buffer :: [boolean()]) ->
			[boolean()].
check_file(FileHandler, PieceLength, ID,
	   Length, PieceChecks, Buffer) when Length =< PieceLength ->
    Piece = file_handler:read(FileHandler, ID * PieceLength,
			      Length),
    Valid = check_piece(Piece, PieceChecks),
    lists:reverse([Valid | Buffer]);
check_file(FileHandler, PieceLength, ID,
	   Length, PieceChecks, Buffer) ->
    ThisCheck = string:substr(PieceChecks, 1, 20),
    RestCheck = string:substr(PieceChecks, 21),
    Piece = file_handler:read(FileHandler, ID * PieceLength,
			      PieceLength),
    Valid = check_piece(Piece, ThisCheck),
    check_file(FileHandler, PieceLength, ID + 1, Length - PieceLength,
	       RestCheck, [Valid | Buffer]).

%%--------------------------------------------------------------------
%% @doc
%% Validate piece
%% @end
%%--------------------------------------------------------------------
-spec check_piece(PieceData :: binary(), PieceCheck :: string()) ->
			 boolean().
check_piece(PieceData, PieceCheck) ->
    PieceCheck == binary_to_list(crypto:sha(PieceData)).

%%--------------------------------------------------------------------
%% @doc
%% Tracker is connected, do things with the peers returned
%% @end
%%--------------------------------------------------------------------
-spec tracker_connection(Peers :: [tracker:address()],
			 State :: #state{}) ->
				NewState :: #state{}.
tracker_connection(Peers, #state{peers = Peers} = State) ->
    %list of peers is the same as we already have, do nothing
    State;
tracker_connection(Peers,
		   #state{peers = OldPeers,
			  torrent = #torrent{info_hash = IH},
			  pieces = Have, supervisor = Sup} = State) ->
    NewPeers = ordsets:subtract(ordsets:from_list(Peers), OldPeers),
    do_things_with_new_peers_here,
    get_piece(Have),
    connect_peers(NewPeers, 5, IH, Sup),
    State.

%%--------------------------------------------------------------------
%% @doc
%% Connect peers
%% @end
%%--------------------------------------------------------------------
-spec connect_peers(Peers :: ordsets:ordset(tracker:address()),
		    MaxStartedConnections :: integer(),
		    InfoHash :: string(), Sup :: pid()) ->
			   [tracker:address()].
connect_peers(_, 0, _, _) ->
    [];
connect_peers([], _, _, _) ->
    [];
connect_peers([Head | Tail], Max, InfoHash, Sup) when Max > 0 ->
    case connect(Head, InfoHash) of
	{error, _Reason} ->
	    connect_peers(Tail, Max, InfoHash, Sup);
	{ok, Socket} ->
	    torrent_sup:start_peer_connections(Sup, Head, Socket),
	    start_send_recv_processes,
	    %gen_tcp:close(Socket),
	    [Head | connect_peers(Tail, Max - 1, InfoHash, Sup)]
    end.

%%--------------------------------------------------------------------
%% @doc
%% Generate an "empty" "have" list
%% @end
%%--------------------------------------------------------------------
-spec generate_empty_have(Length :: integer(),
			  PieceLength :: integer()) -> 
				 [boolean()].
generate_empty_have(Length, PieceLength) when Length < PieceLength ->
    [false | []];
generate_empty_have(Length, PieceLength) ->
    [false | generate_empty_have(Length - PieceLength, PieceLength)].


%%--------------------------------------------------------------------
%% @doc
%% Pick a piece to download
%% @end
%%--------------------------------------------------------------------
-spec get_piece(HavePieces :: [boolean()]) -> integer().
get_piece([false | _Tail]) ->
    0;
get_piece([true | Tail]) ->
    1 + get_piece(Tail).



%%--------------------------------------------------------------------
%% @doc
%% Connect to a peer
%% @end
%%--------------------------------------------------------------------
-spec connect(Peer :: tracker:address(), InfoHash :: string()) ->
		     {ok, Socket :: gen_tcp:socket()} |
		     {error, Reason :: term()}.
connect({Address, Port}, IH) ->
    PeerId = list_to_binary(singleton_storage:get_peerid()),
    InfoHash = list_to_binary(IH),
    case gen_tcp:connect(Address, Port,
			 [{active, false}, binary], 1000) of
	{error, Reason} ->
	    {error, Reason};
	{ok, Sock} ->
	    do_handshake(Sock, PeerId, InfoHash)
    end.

%%--------------------------------------------------------------------
%% @doc
%% Handshake with a peer
%% @end
%%--------------------------------------------------------------------
-spec do_handshake(Socket :: gen_tcp:socket(), PeerID :: binary(),
		   InfoHash :: binary()) ->
			  {ok, Socket :: gen_tcp:socket()} |
			  {error, Reason :: term()}.
do_handshake(Socket, PeerID, InfoHash) ->
    Reserved = <<0:64/big>>,
    Binary = <<19:8/integer,
	       "BitTorrent protocol",
	       Reserved/binary,
	       InfoHash/binary,
	       PeerID/binary>>,
    gen_tcp:send(Socket, Binary),
    case gen_tcp:recv(Socket, 68) of
	{ok, <<19:8/integer, "BitTorrent protocol", _:8/binary,
	       InfoHash:20/binary, _RemoteId:20/binary>>} ->
	    {ok, Socket};
	{ok, _Data} ->
	    {error, invalid_handshake};
	{error, Error} ->
	    {error, Error}
    end.
